from rest_framework import viewsets, status
from rest_framework.decorators import action
from rest_framework.response import Response
from rest_framework.permissions import IsAuthenticated
from django.db.models import Count

from .models import Conversation, Message, QuestionFeedback
from .serializers import (
    ConversationSerializer,
    ConversationDetailSerializer,
    MessageSerializer,
    AskQuestionSerializer,
    QuestionFeedbackSerializer
)
from .services import RAGService


class ConversationViewSet(viewsets.ModelViewSet):
    """对话会话视图集"""
    permission_classes = [IsAuthenticated]

    def get_queryset(self):
        queryset = Conversation.objects.filter(user=self.request.user)

        # 按归档状态筛选
        is_archived = self.request.query_params.get('is_archived')
        if is_archived is not None:
            queryset = queryset.filter(is_archived=is_archived.lower() == 'true')

        # 按知识库筛选
        kb_id = self.request.query_params.get('knowledge_base')
        if kb_id:
            queryset = queryset.filter(knowledge_base_id=kb_id)

        return queryset

    def get_serializer_class(self):
        if self.action == 'retrieve':
            return ConversationDetailSerializer
        return ConversationSerializer

    @action(detail=True, methods=['post'])
    def archive(self, request, pk=None):
        """归档对话"""
        conversation = self.get_object()
        conversation.is_archived = True
        conversation.save()
        return Response({'message': '对话已归档'})

    @action(detail=True, methods=['post'])
    def unarchive(self, request, pk=None):
        """取消归档"""
        conversation = self.get_object()
        conversation.is_archived = False
        conversation.save()
        return Response({'message': '已取消归档'})

    @action(detail=False, methods=['delete'])
    def clear_all(self, request):
        """清空所有对话"""
        count = self.get_queryset().count()
        self.get_queryset().delete()
        return Response({'message': f'已删除 {count} 个对话'})


class MessageViewSet(viewsets.ReadOnlyModelViewSet):
    """消息视图集"""
    serializer_class = MessageSerializer
    permission_classes = [IsAuthenticated]

    def get_queryset(self):
        return Message.objects.filter(conversation__user=self.request.user)


class QAViewSet(viewsets.ViewSet):
    """问答视图集"""
    permission_classes = [IsAuthenticated]

    @action(detail=False, methods=['post'])
    def ask(self, request):
        """提问"""
        serializer = AskQuestionSerializer(data=request.data)
        if not serializer.is_valid():
            return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)

        question = serializer.validated_data['question']
        conversation_id = serializer.validated_data.get('conversation_id')
        kb_id = serializer.validated_data.get('knowledge_base_id')
        use_rag = serializer.validated_data.get('use_rag', True)
        top_k = serializer.validated_data.get('top_k', 3)

        # 获取或创建对话
        if conversation_id:
            try:
                conversation = Conversation.objects.get(id=conversation_id, user=request.user)
            except Conversation.DoesNotExist:
                return Response({'error': '对话不存在'}, status=status.HTTP_404_NOT_FOUND)

            # 如果未明确指定知识库，沿用对话已有的知识库
            if kb_id is None:
                kb_id = conversation.knowledge_base_id
        else:
            # 创建新对话
            conversation = Conversation.objects.create(
                user=request.user,
                title=question[:50],
                knowledge_base_id=kb_id
            )

        # 未选择知识库时走通用问答，关闭RAG检索
        if not kb_id:
            use_rag = False

        # 保存用户问题
        user_message = Message.objects.create(
            conversation=conversation,
            role='user',
            content=question
        )

        # 调用RAG服务生成答案
        rag_service = RAGService()
        try:
            answer_data = rag_service.generate_answer(
                question=question,
                user=request.user,
                knowledge_base_id=kb_id,
                use_rag=use_rag,
                top_k=top_k
            )

            # 保存助手回复
            assistant_message = Message.objects.create(
                conversation=conversation,
                role='assistant',
                content=answer_data['answer'],
                tokens=answer_data.get('tokens'),
                model=answer_data.get('model')
            )

            # 保存检索到的文档
            # for doc_info in answer_data.get('retrieved_docs', []):
            #     RetrievedDocument.objects.create(
            #         message=assistant_message,
            #         document_id=doc_info['document_id'],
            #         chunk_index=doc_info.get('chunk_index'),
            #         relevance_score=doc_info['score'],
            #         content=doc_info['content']
            #     )

            return Response({
                'conversation_id': conversation.id,
                'question': question,
                'answer': answer_data['answer'],
                'retrieved_docs': answer_data.get('retrieved_docs', []),
                'model': answer_data.get('model'),
            })

        except Exception as e:
            return Response({
                'error': f'生成答案失败: {str(e)}'
            }, status=status.HTTP_500_INTERNAL_SERVER_ERROR)

    @action(detail=False, methods=['get'])
    def history(self, request):
        """获取问答历史"""
        conversations = Conversation.objects.filter(user=request.user, is_archived=False)[:10]
        serializer = ConversationSerializer(conversations, many=True)
        return Response(serializer.data)


class FeedbackViewSet(viewsets.ModelViewSet):
    """反馈视图集"""
    serializer_class = QuestionFeedbackSerializer
    permission_classes = [IsAuthenticated]

    def get_queryset(self):
        return QuestionFeedback.objects.filter(user=self.request.user)
